from scipy.io import loadmat, savemat
import scipy.io as scio
from random import random
import numpy as np
import math
from hmmlearn.hmm import GaussianHMM
from sklearn import svm
import warnings
import time

import config
from feature import load_features
from train import train, train_window_partition, svm_test_window_partition, test_window_partition, validate, calc_class_features, simpleSVMtrain


warnings.filterwarnings("ignore")


def test(model1, model2, classifer):
	if config.TEST_TYPE == 2:
		test_file = config.SIMPLE_HMM_TEST_FILE
	else:
		test_file = config.TEST_FILE
		
	beginTime = time.clock()

	for filename in range(len(test_file)):	
		test_dir = config.TEST_DIR + "feature-" + test_file[filename] + ".mat"
		features, segment_num = load_features(test_dir)
		data = test_window_partition(features)

		print("Load train data ok!")
		print("\tPartition with window size = {}, now data size = {}".format(config.WINDOW_K, data.shape))

		features = calc_class_features(model1, model2, data)

		print("Prepare for SVM classifer validation...")
		print("\tFeature size = {}".format(features.shape))

		if config.TEST_TYPE == 2:
			predict_labels = []
			for f in features:
				if f[config.WINDOW_K] < f[config.WINDOW_K*2+1]:
					predict_labels.append(1)
				else:
					predict_labels.append(0)
			len_predict_labels = len(predict_labels)
		else:
			predict_labels = classifer.predict(features)
			len_predict_labels = predict_labels.shape[0]
		#print(predict_labels)
		result = []
		window_num = segment_num-config.WINDOW_K+1
		for i in range(int(math.floor(float(len_predict_labels)/float(window_num)))):
			if sum(predict_labels[i*window_num: (i+1)*window_num]) > 0:
				result.append(1)
			else:
				result.append(0)
		print(result)
		endTime = time.clock()
		if config.TEST_TYPE == 2:
			scio.savemat(config.RESULT_DIR + "HMMresult-" + str(config.WINDOW_K) + "-" + test_file[filename] + ".mat", {'result':result, 'time':endTime-beginTime})
		else:
			scio.savemat(config.RESULT_DIR + "result-" + test_file[filename] + ".mat", {'result':result})

def simpleSVMtest(classifier):
	for filename in range(len(config.TEST_FILE)):	
		test_dir = config.TEST_DIR + "feature-" + config.TEST_FILE[filename] + ".mat"
		features, segment_num = load_features(test_dir)
		data = svm_test_window_partition(features)
		predict_labels = classifier.predict(data)
		result = []
		window_num = segment_num-config.WINDOW_K+1
		for i in range(int(math.floor(float(predict_labels.shape[0])/float(window_num)))):
			if sum(predict_labels[i*window_num: (i+1)*window_num]) > 0:
				result.append(1)
			else:
				result.append(0)
		print(result)
		scio.savemat(config.RESULT_DIR + "SVMresult-" + config.TEST_FILE[filename] + ".mat", {'result':result})

if __name__ == '__main__':
	if config.TEST_TYPE == 0:
		svmClassifier = simpleSVMtrain()
		simpleSVMtest(svmClassifier)
	else:
		model1, model2, classifer = train()
		print("Load models ok!")

		print("Begin test...")
		test(model1, model2, classifer)